#pragma once

#ifndef DIFFIE_HELLMAN
#define DIFFIE_HELLMAN

#include "../Lab1_Large_numbers/__uintX2.h"
#include "../Lab2_Primality_test/solovay_strassen_test.h"
#include "../Lab3_Primality_test/poklington_test.h"


class subscriber
{
public:
	subscriber(uint128 p, uint128 g, char* name=NULL)
		:_name(name),_p(p),_g(g)
	{
		_RK_generate();
	}
#ifdef TEST
	subscriber(uint128 p, uint128 g, uint128 key, char* name=NULL)
		:_name(name),_p(p),_g(g),RK(key){}
#endif /* TEST */
	char* get_name() const
	{
		return _name;
	}
	uint128 send_key() const
	{
		return uint128::pow_m(_g,RK,_p);
	}
	void receive_key(const uint128& key)
	{
		SK=uint128::pow_m(key,RK,_p);
	}
private:
	void _RK_generate() 
	{
		RK = uint128::random()%(_p-3) + 2; // RK=2..p-2
	}
	uint128 _p,_g;
	char* _name;
	uint128 RK;
//public:
	uint128 SK;  
};

class DH_System
{
public:
	DH_System()
	{
		create_gen();
		A = new subscriber(_p,_g,"Alice");
		B = new subscriber(_p,_g,"Bob");
	}
	
	void key_exchange()
	{
		// Alice puts key into medium 
		a_msg = A->send_key();
        // Bob gets key from medium
		B->receive_key(a_msg);
		// Bob puts key into medium 
		b_msg = B->send_key();
        // Alice gets key from medium
		A->receive_key(b_msg);
	}
	subscriber *A,*B;

#ifdef TEST
	DH_System(const uint128 p, 
		      const uint128 g, 
			  const uint128 A_key,
			  const uint128 B_key)
		:_p(p),_g(g)
	{
		A = new subscriber(_p, _g, A_key, "Alice");
		B = new subscriber(_p, _g, B_key, "Bob");
	}
#endif /* TEST */

	//emulate open for passive listening medium
	uint128 get_p() const
	{
		return _p;
	}
	uint128 get_g() const
	{
		return _g;
	}
	uint128 get_a_msg() const
	{
		return a_msg;
	}
	uint128 get_b_msg() const
	{
		return b_msg;
	}

private:
	void create_gen()
	{
		srand(int(time(NULL)));
		_p = poklington().get_prime(sizeof(uint128)*8, _g);
	}
	uint128 _p,_g;
	uint128 a_msg,b_msg;
};

#endif /* DIFFIE_HELLMAN */